package aiaudio.storage.simple;

import aiaudio.storage.api.IArtist;
import aiaudio.storage.api.IAudioTestSet;
import aiaudio.storage.api.IAudioTrainingSet;
import aiaudio.storage.api.IUser;
import aiaudio.storage.api.RevereUserArtistTuple;
import aiaudio.storage.api.Tuple;
import aiaudio.storage.api.UserArtistTuple;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;

/**
 *
 * @author Nastya
 */
class StorageSplitter {

    private List<IUser> users;
    private List<IArtist> artists;
    private HashMap<RevereUserArtistTuple, Double> ratings;

    StorageSplitter(List<IUser> users, List<IArtist> artists,
            HashMap<RevereUserArtistTuple, Double> ratings) {
        this.users = users;
        this.artists = artists;
        this.ratings = ratings;
    }

    IAudioTrainingSet extractTrainingSet(int artistsStart, int artistsEnd,
            int usersStart, int usersEnd) {
        List<IArtist> firstIArtistsSet =
                new ArrayList<IArtist>(artists.subList(artistsStart,
                artistsEnd));
        List<IUser> firstIUsersSet =
                new ArrayList<IUser>(users.subList(usersStart,
                usersEnd));
        HashMap<IUser, List<Tuple<IArtist, Double>>> firstRatingsSet =
                collectTrainingRatingSet(firstIUsersSet, firstIArtistsSet, ratings);
        return new AudioTrainingSet(firstIUsersSet, firstIArtistsSet,
                firstRatingsSet);

    }

    IAudioTestSet extractTestSet(int artistsStart, int artistsEnd,
            int usersStart, int usersEnd) {
        List<IArtist> testIArtistsSet =
                new ArrayList<IArtist>(artists.subList(artistsStart,
                artistsEnd));
        List<IUser> testIUsersSet =
                new ArrayList<IUser>(users.subList(usersStart,
                usersEnd));
        HashMap<UserArtistTuple, Double> testRatingsSet =
                collectTestRatingSet(testIUsersSet, testIArtistsSet, ratings);
        return new AudioTestSet(users, artists, testRatingsSet);

    }

    private HashMap<IUser, List<Tuple<IArtist, Double>>> collectTrainingRatingSet(
            List<IUser> usersSet, List<IArtist> artistsSet,
            HashMap<RevereUserArtistTuple, Double> ratings) {
        HashMap<IUser, List<Tuple<IArtist, Double>>> result =
                new HashMap<IUser, List<Tuple<IArtist, Double>>>();
        for (Entry<RevereUserArtistTuple, Double> entry : ratings.entrySet()) {
            addRatingToTrainingSetIfNeeded(entry, usersSet, artistsSet, result);
        }
        return result;
    }

    private HashMap<UserArtistTuple, Double> collectTestRatingSet(
            List<IUser> usersSet,
            List<IArtist> artistSet,
            HashMap<RevereUserArtistTuple, Double> ratings) {

        HashMap<UserArtistTuple, Double> result =
                new HashMap<UserArtistTuple, Double>();
        for (Entry<RevereUserArtistTuple, Double> entry : ratings.entrySet()) {
            addRatingToTestSetIfNeeded(entry, usersSet, artistSet, result);
        }

        return result;
    }

    private void addRatingToTrainingSetIfNeeded(
            Entry<RevereUserArtistTuple, Double> entry,
            List<IUser> usersSet, List<IArtist> artistsSet,
            HashMap<IUser, List<Tuple<IArtist, Double>>> result) {
        RevereUserArtistTuple userArtistTuple = entry.getKey();
        IArtist artist = userArtistTuple.getFirst();
        IUser user = userArtistTuple.getSecond();
        boolean setsContainsTuple =
                usersSet.contains(user) && artistsSet.contains(artist);

        if (setsContainsTuple) {
            putRatingToTrainingSet(result, user, entry, artist);
        }
    }

    private void putRatingToTrainingSet(
            HashMap<IUser, List<Tuple<IArtist, Double>>> result,
            IUser user, Entry<RevereUserArtistTuple, Double> entry,
            IArtist artist) {
        List<Tuple<IArtist, Double>> userRatings;
        if (result.containsKey(user)) {
            userRatings = result.get(user);
        } else {
            userRatings = new ArrayList<Tuple<IArtist, Double>>();
            result.put(user, userRatings);
        }

        Double rating = entry.getValue();
        userRatings.add(new Tuple<IArtist, Double>(artist, rating));
    }

    private void addRatingToTestSetIfNeeded(
            Entry<RevereUserArtistTuple, Double> entry,
            List<IUser> usersSet, List<IArtist> artistSet,
            HashMap<UserArtistTuple, Double> result) {
        RevereUserArtistTuple userArtistTuple = entry.getKey();

        IArtist artist = userArtistTuple.getFirst();
        IUser user = userArtistTuple.getSecond();
        boolean setsContainsTuple =
                usersSet.contains(user) && artistSet.contains(artist);

        if (setsContainsTuple) {
            UserArtistTuple userArtistRuple =
                    userArtistTuple.cloneReverse();
            Double rating = entry.getValue();
            result.put(userArtistRuple, rating);
        }
    }
}
